from utils import paddle_aux
import paddle


class sa_layer(paddle.nn.Layer):
    """Constructs a Channel Spatial Group module.
    Args:
        k_size: Adaptive selection of kernel size
    """

    def __init__(self, channel, groups=64):
        super(sa_layer, self).__init__()
        self.groups = groups
        self.avg_pool = paddle.nn.AdaptiveAvgPool2D(output_size=1)
        self.cweight = paddle.base.framework.EagerParamBase.from_tensor(tensor
            =paddle.zeros(shape=[1, channel // (2 * groups), 1, 1]))
        self.cbias = paddle.base.framework.EagerParamBase.from_tensor(tensor
            =paddle.ones(shape=[1, channel // (2 * groups), 1, 1]))
        self.sweight = paddle.base.framework.EagerParamBase.from_tensor(tensor
            =paddle.zeros(shape=[1, channel // (2 * groups), 1, 1]))
        self.sbias = paddle.base.framework.EagerParamBase.from_tensor(tensor
            =paddle.ones(shape=[1, channel // (2 * groups), 1, 1]))
        self.sigmoid = paddle.nn.Sigmoid()
        self.gn = paddle.nn.GroupNorm(num_groups=channel // (2 * groups),
            num_channels=channel // (2 * groups))

    @staticmethod
    def channel_shuffle(x, groups):
        b, c, h, w = tuple(x.shape)
        x = x.reshape(b, groups, -1, h, w)
        x = x.transpose(perm=[0, 2, 1, 3, 4])
        x = x.reshape(b, -1, h, w)
        return x

    def forward(self, x):
        b, c, h, w = tuple(x.shape)
        x = x.reshape(b * self.groups, -1, h, w)
        x_0, x_1 = x.chunk(chunks=2, axis=1)
        xn = self.avg_pool(x_0)
        xn = self.cweight * xn + self.cbias
        xn = x_0 * self.sigmoid(xn)
        xs = self.gn(x_1)
        xs = self.sweight * xs + self.sbias
        xs = x_1 * self.sigmoid(xs)
        out = paddle.concat(x=[xn, xs], axis=1)
        out = out.reshape(b, -1, h, w)
        out = self.channel_shuffle(out, 2)
        return out
